#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
重新绘制control_yu.pdf的脚本
解决纵坐标尺度不自适应的问题

使用方法:
python replot_control_yu.py --exp_dir _results/threetanks/67
"""

import numpy as np
import matplotlib.pyplot as plt
import argparse
import os
import sys
from pathlib import Path
import pandas as pd

# 添加项目路径到sys.path
project_root = Path(__file__).parent
sys.path.append(str(project_root))

try:
    from QYtool import datatool
except ImportError:
    print("警告: 无法导入QYtool.datatool，将使用numpy直接加载数据")
    datatool = None

def load_data_files(data_dir):
    """加载保存的数据文件"""
    data = {}
    
    # 需要加载的数据文件列表
    data_files = [
        'ypath', 'upath', 'ypath_var', 'upath_var', 't', 't_step',
        'upath_step', 'upath_var_step', 'us_step', 'ys_step',
        'u_mean', 'u_var', 'y_mean', 'y_var', 'yloss_mean', 'yloss_var',
        'Y_open_mean', 'Y_open_var', 'U_dpc_mean', 'Y_dpc_mean', 
        'U_dpc_var', 'Y_dpc_var'
    ]
    
    for file_name in data_files:
        file_path = os.path.join(data_dir, f'{file_name}.txt')
        if os.path.exists(file_path):
            try:
                data[file_name] = np.loadtxt(file_path)
                print(f"成功加载: {file_name}.txt")
            except Exception as e:
                print(f"加载 {file_name}.txt 失败: {e}")
        else:
            print(f"文件不存在: {file_path}")
    
    return data

def load_config(exp_dir):
    """加载实验配置"""
    import json
    config_path = os.path.join(exp_dir, 'args.json')
    if os.path.exists(config_path):
        with open(config_path, 'r') as f:
            config = json.load(f)
        return config
    else:
        print(f"配置文件不存在: {config_path}")
        return None

def create_improved_plot(data, config, output_dir, test_num=10, args=None):
    """创建改进的控制图表"""
    
    # 从配置中获取参数
    if config:
        Tini = config.get('Tini', 10)
        x0_std_dev = config.get('x0_std_dev', 0.03)
        sampling_period = 0.1  # 默认采样周期
    else:
        Tini = 10
        x0_std_dev = 0.03
        sampling_period = 0.1
    
    # 尝试从数据中重构必要的变量
    if 'ypath' in data and 'upath' in data:
        ypath = data['ypath']
        upath = data['upath']
        ypath_var = data.get('ypath_var', np.zeros_like(ypath))
        upath_var = data.get('upath_var', np.zeros_like(upath))
        
        # 重构时间轴
        if 't' in data:
            t = data['t']
        else:
            t = np.arange(len(ypath)) * sampling_period
            
        if 't_step' in data:
            t_step = data['t_step']
        else:
            t_step = t
            
    else:
        # 如果没有ypath和upath，尝试从其他数据重构
        if 'y_mean' in data and 'u_mean' in data:
            y_mean = data['y_mean']
            u_mean = data['u_mean']
            y_var = data.get('y_var', np.zeros_like(y_mean))
            u_var = data.get('u_var', np.zeros_like(u_mean))
            
            # 假设有初始数据
            yini_mean = np.zeros((Tini, y_mean.shape[1]))
            uini_mean = np.zeros((Tini, u_mean.shape[1]))
            yini_var = np.zeros_like(yini_mean)
            uini_var = np.zeros_like(uini_mean)
            
            ypath = np.vstack((yini_mean, y_mean))
            upath = np.vstack((uini_mean, u_mean))
            ypath_var = np.vstack((yini_var, y_var))
            upath_var = np.vstack((uini_var, u_var))
            
            t = np.arange(len(ypath)) * sampling_period
            t_step = t
        else:
            print("错误: 无法找到足够的数据来绘制图表")
            return False
    
    # 获取其他数据
    ys_step = data.get('ys_step', np.zeros_like(ypath))
    us_step = data.get('us_step', np.zeros_like(upath))
    Y_dpc_mean = data.get('Y_dpc_mean', np.zeros_like(ypath))
    Y_dpc_var = data.get('Y_dpc_var', np.zeros_like(ypath))
    Y_open_mean = data.get('Y_open_mean', None)
    Y_open_var = data.get('Y_open_var', None)
    YLoss_mean = data.get('yloss_mean', np.zeros(len(ypath)))
    
    # 检查数据维度
    if len(ypath.shape) == 1:
        ypath = ypath.reshape(-1, 1)
        ypath_var = ypath_var.reshape(-1, 1)
    if len(upath.shape) == 1:
        upath = upath.reshape(-1, 1)
        upath_var = upath_var.reshape(-1, 1)
    
    n_outputs = ypath.shape[1]
    n_inputs = upath.shape[1]
    
    print(f"数据维度: ypath={ypath.shape}, upath={upath.shape}")
    print(f"输出数量: {n_outputs}, 输入数量: {n_inputs}")
    
    # 创建图表
    fig, ax = plt.subplots(ncols=max(n_outputs, n_inputs), nrows=2, sharex=True, figsize=(16, 8))
    
    # 确保ax是2D数组
    if ax.ndim == 1:
        ax = ax.reshape(2, -1)
    
    tini = Tini * sampling_period
    flag_openloop = Y_open_mean is not None and Y_open_var is not None
    
    # 绘制输出图
    for i in range(n_outputs):
        col_idx = min(i, ax.shape[1] - 1)
        
        # 计算自适应的y轴范围
        y_data_all = []
        y_data_all.extend(ypath[:, i])
        if flag_openloop:
            y_data_all.extend(Y_open_mean[:, i])
        y_data_all.extend(Y_dpc_mean[:, i])
        if len(ys_step) > 0 and ys_step.shape[1] > i:
            y_data_all.extend(ys_step[:, i])
        
        y_min_auto = np.min(y_data_all) - 0.1 * (np.max(y_data_all) - np.min(y_data_all))
        y_max_auto = np.max(y_data_all) + 0.1 * (np.max(y_data_all) - np.min(y_data_all))
        
        # 绘制数据
        mask_before = t <= tini
        mask_after = t >= tini
        
        # PID阶段
        if np.any(mask_before):
            ax[0, col_idx].fill_between(t[mask_before], 
                                      ypath[mask_before, i] - ypath_var[mask_before, i], 
                                      ypath[mask_before, i] + ypath_var[mask_before, i], 
                                      color='lightskyblue', alpha=0.4, label='PID variance')
            ax[0, col_idx].plot(t[mask_before], ypath[mask_before, i], '-', 
                              color='darkviolet', label='PID mean')
        
        # DDeePC阶段
        if np.any(mask_after):
            ax[0, col_idx].fill_between(t[mask_after], 
                                      ypath[mask_after, i] - ypath_var[mask_after, i], 
                                      ypath[mask_after, i] + ypath_var[mask_after, i], 
                                      color='tomato', alpha=0.4, label='DDeePC variance')
            ax[0, col_idx].plot(t[mask_after], ypath[mask_after, i], '-', 
                              color='red', label='DDeePC mean')
        
        # DeePC对比
        if np.any(mask_after) and Y_dpc_mean.shape[1] > i:
            ax[0, col_idx].fill_between(t[mask_after], 
                                      Y_dpc_mean[mask_after, i] - Y_dpc_var[mask_after, i], 
                                      Y_dpc_mean[mask_after, i] + Y_dpc_var[mask_after, i], 
                                      color='greenyellow', alpha=0.4, label='DeePC variance')
            ax[0, col_idx].plot(t[mask_after], Y_dpc_mean[mask_after, i], '-.', 
                              color='dodgerblue', label='DeePC mean')
        
        # 开环对比
        if flag_openloop and Y_open_mean.shape[1] > i:
            ax[0, col_idx].fill_between(t, 
                                      Y_open_mean[:, i] - Y_open_var[:, i], 
                                      Y_open_mean[:, i] + Y_open_var[:, i], 
                                      color='orange', alpha=0.4, label='Open-loop variance')
            ax[0, col_idx].plot(t, Y_open_mean[:, i], '-.', 
                              color='darkorange', label='Open-loop mean')
        
        # 设定点
        if len(ys_step) > 0 and ys_step.shape[1] > i:
            ax[0, col_idx].plot(t_step, ys_step[:, i], linestyle='--', 
                              color='g', label='Set-point')
        
        # Tini分界线
        ax[0, col_idx].axvline(x=tini, linestyle='--', color='b', label='Tini')
        
        # 设置标签和范围
        ax[0, col_idx].set_xlabel('Time (day)')
        ax[0, col_idx].set_ylabel(f"y{i+1}")
        ax[0, col_idx].set_ylim(y_min_auto, y_max_auto)
        
        print(f"输出 y{i+1}: y_range = [{y_min_auto:.4f}, {y_max_auto:.4f}]")
    
    # 绘制输入图
    for i in range(n_inputs):
        col_idx = min(i, ax.shape[1] - 1)
        
        # 计算自适应的u轴范围
        u_data_all = []
        u_data_all.extend(upath[:, i])
        if len(us_step) > 0 and us_step.shape[1] > i:
            u_data_all.extend(us_step[:, i])
        
        u_min_auto = np.min(u_data_all) - 0.1 * (np.max(u_data_all) - np.min(u_data_all))
        u_max_auto = np.max(u_data_all) + 0.1 * (np.max(u_data_all) - np.min(u_data_all))
        
        # 绘制数据
        mask_before = t <= tini
        mask_after = t >= tini
        
        # PID阶段
        if np.any(mask_before):
            ax[1, col_idx].fill_between(t[mask_before], 
                                      upath[mask_before, i] - upath_var[mask_before, i], 
                                      upath[mask_before, i] + upath_var[mask_before, i], 
                                      color='lightskyblue', alpha=0.4, label='PID variance')
            ax[1, col_idx].plot(t[mask_before], upath[mask_before, i], '-', 
                              color='darkviolet', label='PID mean')
        
        # DDeePC阶段
        if np.any(mask_after):
            ax[1, col_idx].fill_between(t[mask_after], 
                                      upath[mask_after, i] - upath_var[mask_after, i], 
                                      upath[mask_after, i] + upath_var[mask_after, i], 
                                      color='orange', alpha=0.4, label='DDeePC variance')
            ax[1, col_idx].plot(t[mask_after], upath[mask_after, i], '-', 
                              color='red', label='DDeePC mean')
        
        # 设定点
        if len(us_step) > 0 and us_step.shape[1] > i:
            ax[1, col_idx].plot(t_step, us_step[:, i], linestyle='--', 
                              color='g', label='Set-point')
        
        # Tini分界线
        ax[1, col_idx].axvline(x=tini, linestyle='--', color='b', label='Tini')
        
        # 设置标签和范围
        ax[1, col_idx].set_xlabel('Time (day)')
        ax[1, col_idx].set_ylabel(f"u{i+1}")
        ax[1, col_idx].set_ylim(u_min_auto, u_max_auto)
        
        print(f"输入 u{i+1}: u_range = [{u_min_auto:.4f}, {u_max_auto:.4f}]")
    
    # 添加图例
    ax[0, 0].legend()
    
    # 添加标题
    avg_loss = np.mean(YLoss_mean) if len(YLoss_mean) > 0 else 0
    fig.suptitle(f"average y loss: {avg_loss:.6f}\n x0 std dev: {x0_std_dev}\n test number: {test_num}")
    
    # 保存图表
    output_path = os.path.join(output_dir, "control_yu_improved.pdf")
    fig.savefig(output_path, bbox_inches='tight', dpi=300)
    plt.close()
    
    print(f"改进的图表已保存到: {output_path}")
    
    # 保存折线图数据到CSV文件（使用与绘图相同的时间轴和数据）
    # 构建参数字典
    csv_args = {
        'sampling_period': sampling_period,
        'Tini': Tini
    }
    save_plot_data_to_csv(data, csv_args, output_dir)
    
    return True

def save_plot_data_to_csv(data, args, output_dir):
    """将绘图数据保存为CSV文件，确保与绘图函数使用相同的时间轴和数据"""
    try:
        # 获取基本参数
        sampling_period = args.get('sampling_period', 0.025)
        Tini = args.get('Tini', 40)
        
        # 重构与绘图函数完全相同的数据
        if 'ypath' in data and 'upath' in data:
            ypath = data['ypath']
            upath = data['upath']
            ypath_var = data.get('ypath_var', np.zeros_like(ypath))
            upath_var = data.get('upath_var', np.zeros_like(upath))
            
            # 重构时间轴（与绘图函数相同）
            if 't' in data:
                t = data['t']
            else:
                t = np.arange(len(ypath)) * sampling_period
                
            if 't_step' in data:
                t_step = data['t_step']
            else:
                t_step = t
                
        else:
            # 如果没有ypath和upath，尝试从其他数据重构（与绘图函数相同）
            if 'y_mean' in data and 'u_mean' in data:
                y_mean = data['y_mean']
                u_mean = data['u_mean']
                y_var = data.get('y_var', np.zeros_like(y_mean))
                u_var = data.get('u_var', np.zeros_like(u_mean))
                
                # 假设有初始数据
                yini_mean = np.zeros((Tini, y_mean.shape[1]))
                uini_mean = np.zeros((Tini, u_mean.shape[1]))
                yini_var = np.zeros_like(yini_mean)
                uini_var = np.zeros_like(uini_mean)
                
                ypath = np.vstack((yini_mean, y_mean))
                upath = np.vstack((uini_mean, u_mean))
                ypath_var = np.vstack((yini_var, y_var))
                upath_var = np.vstack((uini_var, u_var))
                
                t = np.arange(len(ypath)) * sampling_period
                t_step = t
            else:
                print("错误: 无法找到足够的数据来导出CSV")
                return
        
        # 获取其他数据（与绘图函数相同）
        ys_step = data.get('ys_step', np.zeros_like(ypath))
        us_step = data.get('us_step', np.zeros_like(upath))
        Y_dpc_mean = data.get('Y_dpc_mean', np.zeros_like(ypath))
        Y_dpc_var = data.get('Y_dpc_var', np.zeros_like(ypath))
        Y_open_mean = data.get('Y_open_mean', None)
        Y_open_var = data.get('Y_open_var', None)
        
        # 检查数据维度（与绘图函数相同）
        if len(ypath.shape) == 1:
            ypath = ypath.reshape(-1, 1)
            ypath_var = ypath_var.reshape(-1, 1)
        if len(upath.shape) == 1:
            upath = upath.reshape(-1, 1)
            upath_var = upath_var.reshape(-1, 1)
        
        n_outputs = ypath.shape[1]
        n_inputs = upath.shape[1]
        
        print(f"CSV导出 - 数据维度: ypath={ypath.shape}, upath={upath.shape}")
        print(f"CSV导出 - 时间轴长度: t={len(t)}, t_step={len(t_step)}")
        
        tini = Tini * sampling_period
        flag_openloop = Y_open_mean is not None and Y_open_var is not None
        
        # 处理每个输出
        for i in range(n_outputs):
            output_data = {
                'Time': t  # 使用与绘图相同的时间轴
            }
            
            # DDeePC数据（完整的ypath，包含PID+DDeePC阶段）
            output_data['DDeePC'] = ypath[:, i]
            
            # DeePC数据（只在Tini之后有效，与绘图函数逻辑相同）
            mask_after = t >= tini
            deepc_data = np.full(len(t), np.nan)  # 使用NaN填充无效区域
            if np.any(mask_after) and Y_dpc_mean.shape[1] > i:
                deepc_data[mask_after] = Y_dpc_mean[mask_after, i]
            output_data['DeePC'] = deepc_data
            
            # 开环数据
            if flag_openloop and Y_open_mean.shape[1] > i:
                output_data['Open_loop'] = Y_open_mean[:, i]
            else:
                output_data['Open_loop'] = np.full(len(t), np.nan)
            
            # 设定点数据
            if len(ys_step) > 0 and ys_step.shape[1] > i:
                # 使用t_step对应的数据
                if len(t_step) == len(ys_step):
                    # 将t_step上的数据插值到t上
                    setpoint_interp = np.interp(t, t_step, ys_step[:, i])
                    output_data['Setpoint'] = setpoint_interp
                else:
                    output_data['Setpoint'] = np.full(len(t), np.nan)
            else:
                output_data['Setpoint'] = np.full(len(t), np.nan)
            
            # 验证所有数组长度
            lengths = {key: len(value) for key, value in output_data.items()}
            print(f"y{i+1}数据长度: {lengths}")
            
            if len(set(lengths.values())) > 1:
                print(f"错误: y{i+1}数据长度不一致: {lengths}")
                continue
            
            # 保存为CSV
            df = pd.DataFrame(output_data)
            csv_path = os.path.join(output_dir, f'y{i+1}_data.csv')
            df.to_csv(csv_path, index=False)
            print(f"已保存 y{i+1} 数据到: {csv_path}")
        
        # 处理每个输入
        for i in range(n_inputs):
            input_data = {
                'Time': t  # 使用与绘图相同的时间轴
            }
            
            # DDeePC数据（完整的upath，包含PID+DDeePC阶段）
            input_data['DDeePC'] = upath[:, i]
            
            # 设定点数据
            if len(us_step) > 0 and us_step.shape[1] > i:
                # 使用t_step对应的数据
                if len(t_step) == len(us_step):
                    # 将t_step上的数据插值到t上
                    setpoint_interp = np.interp(t, t_step, us_step[:, i])
                    input_data['Setpoint'] = setpoint_interp
                else:
                    input_data['Setpoint'] = np.full(len(t), np.nan)
            else:
                input_data['Setpoint'] = np.full(len(t), np.nan)
            
            # 验证所有数组长度
            lengths = {key: len(value) for key, value in input_data.items()}
            print(f"u{i+1}数据长度: {lengths}")
            
            if len(set(lengths.values())) > 1:
                print(f"错误: u{i+1}数据长度不一致: {lengths}")
                continue
            
            # 保存为CSV
            df = pd.DataFrame(input_data)
            csv_path = os.path.join(output_dir, f'u{i+1}_data.csv')
            df.to_csv(csv_path, index=False)
            print(f"已保存 u{i+1} 数据到: {csv_path}")
            
    except Exception as e:
        print(f"保存CSV数据时出错: {e}")
        import traceback
        traceback.print_exc()

def main():
    parser = argparse.ArgumentParser(description='重新绘制control_yu.pdf')
    parser.add_argument('--exp_dir', type=str, required=True, 
                       help='实验目录路径，例如: _results/threetanks/67')
    parser.add_argument('--test_num', type=int, default=10, 
                       help='测试数量')
    
    args = parser.parse_args()
    
    # 检查实验目录
    if not os.path.exists(args.exp_dir):
        print(f"错误: 实验目录不存在: {args.exp_dir}")
        return
    
    # 查找数据目录
    control_data_dir = None
    for subdir in ['con_opt', 'no_con_opt']:
        potential_dir = os.path.join(args.exp_dir, 'control', subdir)
        if os.path.exists(potential_dir) and os.listdir(potential_dir):
            control_data_dir = potential_dir
            break
    
    if control_data_dir is None:
        print(f"错误: 在 {args.exp_dir}/control 中找不到数据文件")
        print("请确保已经运行过控制测试并生成了数据")
        return
    
    print(f"使用数据目录: {control_data_dir}")
    
    # 加载数据
    data = load_data_files(control_data_dir)
    if not data:
        print("错误: 无法加载任何数据文件")
        return
    
    # 加载配置
    config = load_config(args.exp_dir)
    
    # 创建输出目录
    output_dir = os.path.join(args.exp_dir, 'control_fig')
    os.makedirs(output_dir, exist_ok=True)
    
    # 创建改进的图表
    success = create_improved_plot(data, config, output_dir, args.test_num, args)
    
    if success:
        print("\n✅ 重新绘图完成!")
        print(f"改进的图表保存在: {output_dir}/control_yu_improved.pdf")
        print("\n主要改进:")
        print("- 自适应纵坐标范围")
        print("- 更好的数据可视化")
        print("- 避免直线显示问题")
        print("- 折线图数据已导出为CSV文件")
    else:
        print("❌ 重新绘图失败")

if __name__ == '__main__':
    main()